package com.sicau.app;

import com.sicau.advisor.LoggerAdvisor;
import com.sicau.advisor.SafeGuardAdvisor;
import com.sicau.chatmemory.BaseOnFileChatMemoryRepository;
import com.sicau.rag.RewriteUserMessageQueryTransformer;
import com.sicau.rag.YouDaoTranslationQueryTransformer;
import com.sicau.utils.YouDaoTranslationFactory;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.List;

@Slf4j
@Component
public class PsychologicalConsultationApp {

    private final VectorStore vectorStore;
    private final RewriteUserMessageQueryTransformer queryTransformer;
    private final YouDaoTranslationQueryTransformer translationQueryTransformer;
    private final ToolCallback[] allTools;
    private final ToolCallbackProvider toolCallbackProvider;

    private final String systemPrompt = """
            你是温暖专业的 AI 心理咨询顾问，仅提供心理支持与状态评估，不替代专业诊疗。遵循以下规则：
            1、共情接纳：不用 “你应该”“别想太多”，不评判用户感受，用 “我能理解”“这种心情很正常” 回应。
            2、引导式提问：围绕情绪（最近心情、是否有低落 / 焦虑等）、睡眠饮食、社交状态、近期压力 / 困扰，用通俗语言逐步了解，单次 1 个问题，避免过载。
            3、安全边界：提及自伤 / 自杀 / 伤害他人，立即提示 “联系信任的人、本地心理援助热线或就医”，不追问细节；不做病理诊断、不开治疗方案。
            4、简洁回应：单次回复 3 句话内，让用户充分表达，动态跟进状态变化。
            """;
    private final ChatClient chatClient;

    public PsychologicalConsultationApp(VectorStore vectorStore, RewriteUserMessageQueryTransformer queryTransformer, YouDaoTranslationFactory youDaoTranslationFactory, ToolCallback[] allTools, ToolCallbackProvider toolCallbackProvider, ChatModel chatModel) {
        this.vectorStore = vectorStore;
        this.queryTransformer = queryTransformer;
        this.translationQueryTransformer = new YouDaoTranslationQueryTransformer(youDaoTranslationFactory);
        this.allTools = allTools;
        this.toolCallbackProvider = toolCallbackProvider;

        ChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(new BaseOnFileChatMemoryRepository("./src/main/resources/temp/chatRepository"))
                .maxMessages(10)
                .build();
        chatClient = ChatClient.builder(chatModel)
                .defaultAdvisors(
                        // 日志Advisor
                        new LoggerAdvisor(),
                        // 会话记忆Advisor
                        MessageChatMemoryAdvisor.builder(chatMemory)
                                .build()
                )
                .build();
    }

    /**
     * 支持多轮对话记忆的ai心理咨询咨询
     * @param message 用户提问
     * @param chatId 会话id
     * @return
     */
    public String doChat(String message, String chatId) {
        String content = chatClient.prompt()
                .user(message)
                .system(systemPrompt)
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, chatId))
                .advisors(// 敏感词Advisor
                        new SafeGuardAdvisor()
                )
                .call()
                .content();
        return content;
    }

    /**
     * 支持多轮对话记忆流式输出的的ai心理咨询咨询
     * @param message 用户提问
     * @param chatId 会话id
     * @return
     */
    public Flux<String> doChatWithSSE(String message, String chatId) {
        Flux<String> content = chatClient.prompt()
                .user(message)
                .system(systemPrompt)
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, chatId))
                .advisors(// 敏感词Advisor
                        new SafeGuardAdvisor()
                )
                .stream()
                .content();
        return content;
    }

    // 心理咨询报告
    record LoveReport(String title, List<String> suggestions) {
    }

    /**
     * 生成心理咨询报告的会话
     * @param message 用户提问
     * @param chatId 会话id
     * @return
     */
    public LoveReport doChatWithReport(String message, String chatId) {
        LoveReport content = chatClient.prompt()
                .user(message)
                .system(systemPrompt + "\n每次对话结束后都要为用户生成一份用户报告，报告标题为{用户名}的心理咨询报告，报告内容为建议列表")
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, chatId))
                .advisors(// 敏感词Advisor
                        new SafeGuardAdvisor("""
                                {
                                    "title": "敏感词异常",
                                    "suggestions": ["请去除敏感词后再次尝试！"]
                                }
                                """)
                )
                .call()
                .entity(LoveReport.class);
        log.info("chat result: {}", content);
        return content;
    }

    /**
     * 聊天+rag
     * @param message
     * @param chatId
     * @return
     */
    public String doChatWithRag(String message, String chatId) {
        // 优化用户提示词
        // message = queryTransformer.rewrite(message);
        message = translationQueryTransformer.transform(message);
        return chatClient.prompt()
                .user(message)
                .system(systemPrompt)
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, chatId))
                .advisors(
                        new SafeGuardAdvisor()
                        // RAGAdvisorFactory.createRetrievalAugmentationAdvisor(vectorStore)
                )
                .call()
                .content();
    }

    /**
     * 工具调用
     * @param message 用户提问
     * @param chatId 会话id
     * @return
     */
    public String doChatWithTool(String message, String chatId) {
        String content = chatClient.prompt()
                .user(message)
                .system(systemPrompt)
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, chatId))
                .advisors(// 敏感词Advisor
                        new SafeGuardAdvisor()
                )
                .toolCallbacks(allTools)
                .call()
                .content();
        return content;
    }

    /**
     * mcp
     * @param message 用户提问
     * @param chatId 会话id
     * @return
     */
    public String doChatWithMcp(String message, String chatId) {
        String content = chatClient.prompt()
                .user(message)
                .system(systemPrompt)
                .advisors(advisorSpec -> advisorSpec.param(ChatMemory.CONVERSATION_ID, chatId))
                .advisors(// 敏感词Advisor
                        new SafeGuardAdvisor()
                )
                .toolCallbacks(toolCallbackProvider)
                .call()
                .content();
        return content;
    }
}
